# -*- coding: utf-8 -*-
"""
普通哈希算法和一致性哈希算法的比较
"""
from typing import Sequence, Dict, Tuple
from hashlib import md5
from struct import unpack_from
from bisect import bisect_left

ITEMS = 10000000
NODES = 100
NEW_NODES = 101


def normal_hash():
    """普通的哈希函数，可以均匀的将数据分布到不同的节点，但是如果增加和删除节点，会导致大量数据迁移"""
    node_stat = [0 for i in range(NODES)]
    for item in range(ITEMS):
        key = md5(str(item).encode()).digest()
        hash_value = unpack_from(">I", key)[0]
        node_id = hash_value % NODES
        node_stat[node_id] += 1

    average = ITEMS / NODES
    max_value = max(node_stat)
    min_value = min(node_stat)

    print(f"Avg: {average}")
    print(f"Max: {max_value}({(max_value - average) * 100.0 / average}%)")
    print(f"Min: {min_value}({(average - min_value) * 100.0 / average}%)")


def normal_hash_add():
    """普通哈希函数在增加一个节点的情况下的表现"""
    node_stat = [0 for i in range(NODES)]
    change = 0
    for item in range(ITEMS):
        key = md5(str(item).encode()).digest()
        hash_value = unpack_from(">I", key)[0]
        node_id = hash_value % NODES
        node_id_new = hash_value % NEW_NODES
        if node_id != node_id_new:
            change += 1

    print(f"Change: {change * 100.0 / ITEMS}%")


def _hash(value):
    key = md5(str(value).encode()).digest()
    hash_value = unpack_from(">I", key)[0]
    return hash_value


def consist_hash():
    """
    普通的一致性hash算法，因为每个节点在环上占空间不一致，数据分布不均匀
    :return:
    """
    ring = []
    hash2node = {}
    node_stat = [0 for i in range(NODES)]

    for n in range(NODES):
        hash_value = _hash(n)
        ring.append(hash_value)
        hash2node[hash_value] = n
    ring.sort()

    for item in range(ITEMS):
        hash_value = _hash(item)
        ring_index = bisect_left(ring, hash_value) % NODES
        node_stat[hash2node[ring[ring_index]]] += 1

    average = ITEMS / NODES
    max_value = max(node_stat)
    min_value = min(node_stat)

    print(f"Avg: {average}")
    print(f"Max: {max_value}({(max_value - average) * 100.0 / average}%)")
    print(f"Min: {min_value}({(average - min_value) * 100.0 / average}%)")


def consist_hash_add():
    """一致性哈希函数，增加一个节点的情况下，需要迁移的数据量
    :return:
    """

    ring = []
    new_ring = []

    for n in range(NODES):
        ring.append(_hash(n))
    ring.sort()

    for n in range(NEW_NODES):
        new_ring.append(_hash(n))
    new_ring.sort()

    change = 0

    for item in range(ITEMS):
        hash_value = _hash(item)
        ring_index = ring[bisect_left(ring, hash_value) % NODES]
        new_ring_index = new_ring[bisect_left(new_ring, hash_value) % NEW_NODES]
        if ring_index != new_ring_index:
            change += 1

    print(f"Change: {change * 100.0 / ITEMS}%")


def virtual_consist_hash(vnodes=1000):
    """添加了虚拟节点的一致性哈希算法，解决数据分配不均"""
    node_stat = [0 for i in range(NODES)]

    ring = []
    hash2node = {}

    for n in range(NODES):
        for v in range(vnodes):
            hash_value = _hash(str(n) + str(v))
            ring.append(hash_value)
            hash2node[hash_value] = n
    ring.sort()

    all_virtual_nodes = NODES * vnodes

    for item in range(ITEMS):
        hash_value = _hash(item)
        ring_index = bisect_left(ring, hash_value) % all_virtual_nodes
        node_stat[hash2node[ring[ring_index]]] += 1

    print(sum(node_stat), node_stat)

    average = ITEMS / NODES
    max_value = max(node_stat)
    min_value = min(node_stat)

    print(f"Avg: {average}")
    print(f"Max: {max_value}({(max_value - average) * 100.0 / average}%)")
    print(f"Min: {min_value}({(average - min_value) * 100.0 / average}%)")


def virtual_consist_hash_add(vnodes=1000):
    """添加了虚拟节点的一致性哈希算法，添加一个节点"""
    ring = []
    hash2node = {}
    ring2 = []
    hash2node2 = {}

    for n in range(NODES):
        for v in range(vnodes):
            hash_value = _hash(str(n) + str(v))
            ring.append(hash_value)
            hash2node[hash_value] = n
    ring.sort()

    for n in range(NEW_NODES):
        for v in range(vnodes):
            hash_value = _hash(str(n) + str(v))
            ring2.append(hash_value)
            hash2node2[hash_value] = n

    ring2.sort()
    print(len(ring), len(ring2))

    all_virtual_nodes = NODES * vnodes
    all_virtual_nodes2 = NEW_NODES * vnodes

    change = 0
    for item in range(ITEMS):
        hash_value = _hash(item)
        ring_index = bisect_left(ring, hash_value) % all_virtual_nodes
        ring_index2 = bisect_left(ring2, hash_value) % all_virtual_nodes2

        if hash2node[ring[ring_index]] != hash2node2[ring2[ring_index2]]:
            change += 1

    print(f"Change: {change * 100.0 / ITEMS}%")


def part_consist_hash(log_node=7, max_power=32):
    """改进一致性哈希算法，使数据的分布更加均匀"""
    partition = max_power - log_node

    node_stat = [0 for i in range(NODES)]

    ring = []
    part2node = {}

    for part in range(2 ** log_node):
        ring.append(part)
        part2node[part] = part % NODES

    for item in range(ITEMS):
        hash_value = _hash(item) >> partition
        part = bisect_left(ring, hash_value)
        node_id = part % NODES
        node_stat[node_id] += 1

    print(f"length of ring: {len(ring)}")

    average = ITEMS / NODES
    max_value = max(node_stat)
    min_value = min(node_stat)

    print(f"Avg: {average}")
    print(f"Max: {max_value}({(max_value - average) * 100.0 / average}%)")
    print(f"Min: {min_value}({(average - min_value) * 100.0 / average}%)")


def part_consist_hash_add(log_node=7, max_power=32):
    """添加一个节点的情况"""
    partition = max_power - log_node

    node_stat = [0 for i in range(NODES)]

    ring = []
    part2node = {}
    part2node2 = {}

    for part in range(2 ** log_node):
        ring.append(part)
        part2node[part] = part % NODES
        part2node2[part] = part % NEW_NODES

    change = 0

    for item in range(ITEMS):
        hash_value = _hash(item) >> partition
        part = bisect_left(ring, hash_value)
        part2 = bisect_left(ring, hash_value)

        node_id = part2node[part] % NODES
        node_id2 = part2node2[part] % NEW_NODES

        if node_id != node_id2:
            change += 1

    print(f"Change: {change * 100.0 / ITEMS}%")


def select_node(key: str, node_range: Dict[str, Tuple[int]], max_range_value=65535) -> str:
    """基于一致性哈希算法选择节点
    :param key: 数据的key值
    :param node_range: 每个节点对应的虚拟节点的范围
    :param max_range_value: range的最大值，总计有max_range_value + 1 个虚拟节点
    :return node id
    """
    hash_value = _hash(key) % max_range_value

    for node_id, vnodes_range in node_range.items():
        if vnodes_range[0] <= hash_value <= vnodes_range[1]:
            return node_id


if __name__ == "__main__":
    # consist_hash()
    # virtual_consist_hash()
    # virtual_consist_hash_add()
    # part_consist_hash()
    # part_consist_hash_add()

    node_range = {"1": [0, 10922],
                  "2": [10923, 21845],
                  "3": [21846, 32767],
                  "4": [32768, 43690],
                  "5": [43691, 54613],
                  "6": [54614, 65534]}

    node_stat = {key: 0 for key in node_range.keys()}

    for key in range(100000):
        node_id = select_node(str(key),node_range)
        node_stat[node_id] += 1

    print(node_stat)

    pass
